import json
from ultralytics import YOLO
from functools import cmp_to_key
import os
import cv2
def A(img_list, model):

    CLS_ID_NAME_MAP = {
        0: 'student_id',  # 学生证
        1: 'subjective_problem',  # 主观问题
        2: 'fillin_problem',  # 填写问题
        3: 'objective_problem'  # 客观问题
    }
    results = model.predict(source=img_list, save=False, imgsz=640)
    data_id = 0
    data = {}
    def compare_items(item1, item2):
        x1_1, y1_1, x2_1, y2_1 = item1['xyxy']
        x1_2, y1_2, x2_2, y2_2 = item2['xyxy']

        if abs(y1_1-y1_2)<50:
            return x1_1-x1_2
        else:
            return y1_1-y1_2
    for result in results:
        id = 0
        res = []
        for box in result.boxes:
            cls_id = box.cls.cpu().numpy()[0]
            x1, y1, x2, y2 = box.xyxy.cpu().numpy()[0]
            x1 = int(x1)
            y1 = int(y1)
            x2 = int(x2)
            y2 = int(y2)
            cls_name = CLS_ID_NAME_MAP[cls_id]

            if cls_id==3 or cls_id==0:
                continue
            res.append({"id": id, "type": cls_name, "xyxy": [x1, y1, x2, y2]})
            id += 1
        res = sorted(res, key=cmp_to_key(compare_items))
        for i in range(len(res)):
            res[i]["id"]=i
        data[data_id]=[x for x in res]
        res.clear()
        data_id += 1
    return data

if __name__ == '__main__':
    json_file = "text.json"

    folder = 'testdata'
    file_names = os.listdir(folder)

    imgs = []
    for file_name in file_names[:10]:
        img_path = os.path.join(folder, file_name)
        img = cv2.imread(img_path)
        if img is None:
            print('无法读取图片')
            exit(-1)
        imgs += [img]
    model = YOLO(model='best.pt')
    res = A(imgs,model)
    save_js = open(json_file,'w')
    json.dump(res,save_js)
    for i in range(0, len(imgs)):
        img = imgs[i]
        data = res[i]
        for x in data:
            id = x["id"]
            x1,y1,x2,y2 = x["xyxy"]
            cv2.rectangle(img, (x1, y1), (x2, y2), (255, 0, 255), 2)
            cv2.putText(img, str(id), (x1, y1 + 5), cv2.FONT_HERSHEY_SIMPLEX, 1, (255, 0, 255), 2)
        cv2.imwrite('debug/'+str(i)+'.jpg',img)


